#!/usr/bin/env python3
"""
Author: Logan.Li
Gitee: https://gitee.com/attacker
email: admin@attacker.club
Date: 2025-01-05 12:25:52
LastEditTime: 2025-04-09
Description: 导出Kubernetes Deployment+Service和Ingress资源
"""

import os
import yaml
import subprocess
import sys

# 尝试导入kubernetes模块，如果失败则尝试安装
try:
    from kubernetes import client, config
    import urllib3
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
except ImportError:
    print("kubernetes模块未安装，尝试安装...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "kubernetes"])
    from kubernetes import client, config
    import urllib3
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def sanitize_resource(obj, resource_type):
    """清理资源对象，移除不必要的字段"""
    # 获取序列化后的对象
    obj_yaml = api_client.sanitize_for_serialization(obj)
    
    # 确保正确的apiVersion和kind在最前面
    if isinstance(obj, client.V1Service):
        obj_yaml['apiVersion'] = 'v1'
        obj_yaml['kind'] = 'Service'
    elif isinstance(obj, client.V1Deployment):
        obj_yaml['apiVersion'] = 'apps/v1'
        obj_yaml['kind'] = 'Deployment'
    elif isinstance(obj, client.V1Ingress):
        obj_yaml['apiVersion'] = 'networking.k8s.io/v1'
        obj_yaml['kind'] = 'Ingress'
    
    # 重新组织字段顺序，确保apiVersion和kind在最前面
    ordered_obj = {
        'apiVersion': obj_yaml.pop('apiVersion'),
        'kind': obj_yaml.pop('kind'),
        'metadata': obj_yaml.pop('metadata', {}),
        'spec': obj_yaml.pop('spec', {})
    }
    
    # 清理metadata中的不必要字段
    metadata_fields = ['selfLink', 'generation', 'creationTimestamp', 
                      'resourceVersion', 'uid', 'managedFields']
    for field in metadata_fields:
        ordered_obj['metadata'].pop(field, None)
    
    # 清理annotations
    annotations_to_remove = [
        'deployment.kubernetes.io/revision',
        'kubectl.kubernetes.io/restartedAt',
        'kubectl.kubernetes.io/last-applied-configuration'
    ]
    
    if 'annotations' in ordered_obj['metadata']:
        for ann in annotations_to_remove:
            ordered_obj['metadata']['annotations'].pop(ann, None)
        
        # 如果annotations为空，则删除该字段
        if not ordered_obj['metadata']['annotations']:
            del ordered_obj['metadata']['annotations']
    
    # 清理spec中的不必要字段
    if resource_type == 'Service':
        if 'clusterIPs' in ordered_obj['spec']:
            del ordered_obj['spec']['clusterIPs']
        if 'clusterIP' in ordered_obj['spec']:
            del ordered_obj['spec']['clusterIP']
        if 'ipFamilies' in ordered_obj['spec']:
            del ordered_obj['spec']['ipFamilies']
        if 'internalTrafficPolicy' in ordered_obj['spec']:
            del ordered_obj['spec']['internalTrafficPolicy']
        if 'ipFamilyPolicy' in ordered_obj['spec']:
            del ordered_obj['spec']['ipFamilyPolicy']
    
    elif resource_type == 'Deployment':
        # 清理Deployment级别的默认字段
        deployment_defaults = ['revisionHistoryLimit', 'progressDeadlineSeconds']
        for field in deployment_defaults:
            if field in ordered_obj['spec']:
                # revisionHistoryLimit默认值是10，progressDeadlineSeconds默认值是600
                if ((field == 'revisionHistoryLimit' and ordered_obj['spec'][field] == 10) or
                    (field == 'progressDeadlineSeconds' and ordered_obj['spec'][field] == 600)):
                    del ordered_obj['spec'][field]
        
        # 清理strategy字段
        if 'strategy' in ordered_obj['spec']:
            strategy = ordered_obj['spec']['strategy']
            # 如果strategy只有type且为RollingUpdate（默认值），则删除整个strategy
            if (len(strategy) == 1 and 
                strategy.get('type') == 'RollingUpdate'):
                del ordered_obj['spec']['strategy']
            else:
                # 保留strategy但移除不需要的字段
                if 'rollingUpdate' in strategy:
                    rolling_update = strategy['rollingUpdate']
                    if 'maxUnavailable' in rolling_update:
                        del rolling_update['maxUnavailable']
                    if 'maxSurge' in rolling_update:
                        del rolling_update['maxSurge']
                    # 如果rollingUpdate为空，则删除
                    if not rolling_update:
                        del strategy['rollingUpdate']
        
        # 清理Pod模板中的默认字段
        if 'template' in ordered_obj['spec']:
            # 清理Pod模板的annotations
            if ('metadata' in ordered_obj['spec']['template'] and 
                'annotations' in ordered_obj['spec']['template']['metadata']):
                pod_annotations = ordered_obj['spec']['template']['metadata']['annotations']
                for ann in annotations_to_remove:
                    pod_annotations.pop(ann, None)
                
                # 如果Pod模板的annotations为空，则删除该字段
                if not pod_annotations:
                    del ordered_obj['spec']['template']['metadata']['annotations']
            
            # 清理Pod spec
            if 'spec' in ordered_obj['spec']['template']:
                pod_spec = ordered_obj['spec']['template']['spec']
                
                # 清理Pod级别的默认字段
                pod_defaults = {
                    'dnsPolicy': 'ClusterFirst',
                    'restartPolicy': 'Always', 
                    'schedulerName': 'default-scheduler',
                    'terminationGracePeriodSeconds': 30
                }
                
                for field, default_value in pod_defaults.items():
                    if field in pod_spec and pod_spec[field] == default_value:
                        del pod_spec[field]
                
                # 清理空的securityContext
                if 'securityContext' in pod_spec and not pod_spec['securityContext']:
                    del pod_spec['securityContext']
                
                # 清理容器中的默认字段
                if 'containers' in pod_spec:
                    for container in pod_spec['containers']:
                        container_defaults = {
                            'terminationMessagePath': '/dev/termination-log',
                            'terminationMessagePolicy': 'File'
                        }
                        
                        for field, default_value in container_defaults.items():
                            if field in container and container[field] == default_value:
                                del container[field]
    
    # 移除status字段和其他不需要的字段
    if 'status' in obj_yaml:
        del obj_yaml['status']
    
    # 添加其他可能需要的字段
    for key, value in obj_yaml.items():
        if key not in ['apiVersion', 'kind', 'metadata', 'spec'] and value is not None:
            ordered_obj[key] = value
    
    return ordered_obj

def export_resources():
    """导出Kubernetes资源到不同目录"""
    # 加载kubeconfig
    config.load_kube_config()
    global api_client
    api_client = client.ApiClient()
    
    # 创建API客户端
    apps_api = client.AppsV1Api(api_client)
    core_api = client.CoreV1Api(api_client)
    networking_api = client.NetworkingV1Api(api_client)
    
    # 获取所有资源
    all_deployments = apps_api.list_deployment_for_all_namespaces().items
    all_services = core_api.list_service_for_all_namespaces().items
    all_ingresses = networking_api.list_ingress_for_all_namespaces().items
    
    # 创建输出目录
    deployments_dir = os.path.join(os.getcwd(), 'deployments')
    ingress_dir = os.path.join(os.getcwd(), 'ingress')
    os.makedirs(deployments_dir, exist_ok=True)
    os.makedirs(ingress_dir, exist_ok=True)
    
    # 按命名空间和组织资源
    namespace_resources = {}
    
    # 处理Deployment
    for dep in all_deployments:
        namespace = dep.metadata.namespace
        dep_name = dep.metadata.name
        
        if namespace not in namespace_resources:
            namespace_resources[namespace] = {'deployments': {}, 'services': {}, 'ingresses': {}}
        
        # 清理并添加Deployment
        cleaned_dep = sanitize_resource(dep, 'Deployment')
        namespace_resources[namespace]['deployments'][dep_name] = cleaned_dep
    
    # 处理所有Service（避免重复处理）
    for svc in all_services:
        namespace = svc.metadata.namespace
        if namespace not in namespace_resources:
            namespace_resources[namespace] = {'deployments': {}, 'services': {}, 'ingresses': {}}
        
        # 清理并添加Service
        cleaned_svc = sanitize_resource(svc, 'Service')
        namespace_resources[namespace]['services'][svc.metadata.name] = cleaned_svc
    
    # 处理Ingress
    for ingress in all_ingresses:
        namespace = ingress.metadata.namespace
        if namespace not in namespace_resources:
            namespace_resources[namespace] = {'deployments': {}, 'services': {}, 'ingresses': {}}
        
        # 清理并添加Ingress
        cleaned_ingress = sanitize_resource(ingress, 'Ingress')
        namespace_resources[namespace]['ingresses'][ingress.metadata.name] = cleaned_ingress
    
    # 导出Deployment+Service到deployments目录
    for namespace, resources in namespace_resources.items():
        # 在deployments目录下创建命名空间子目录
        ns_deployments_dir = os.path.join(deployments_dir, namespace)
        os.makedirs(ns_deployments_dir, exist_ok=True)
        
        # 记录已导出的Service，避免重复
        exported_services = set()
        
        # 导出每个Deployment及其相关Service
        for dep_name, deployment in resources['deployments'].items():
            # 获取Deployment的Pod标签
            pod_labels = deployment.get('spec', {}).get('template', {}).get('metadata', {}).get('labels', {})
            
            # 查找与此Deployment相关的Service
            related_services = []
            
            # 方法1：通过selector匹配（最准确）
            for svc_name, service in resources['services'].items():
                if svc_name in exported_services:
                    continue
                    
                svc_selector = service.get('spec', {}).get('selector', {})
                if svc_selector:
                    # 检查Service的selector是否匹配Deployment的Pod标签
                    if all(pod_labels.get(k) == v for k, v in svc_selector.items()):
                        related_services.append(service)
                        exported_services.add(svc_name)
            
            # 方法2：通过名称精确匹配（作为补充）
            for svc_name, service in resources['services'].items():
                if svc_name in exported_services:
                    continue
                    
                # 只有在Service名称完全等于Deployment名称时才匹配
                # 避免 bos-admin 匹配到 bos-admin-web 这种情况
                if svc_name == dep_name:
                    related_services.append(service)
                    exported_services.add(svc_name)
            
            # 组合Deployment和相关Service
            combined_resources = [deployment]
            combined_resources.extend(related_services)
            
            # 生成文件名
            filename = f"{dep_name}.yaml"
            filepath = os.path.join(ns_deployments_dir, filename)
            
            # 写入文件，确保正确的YAML格式
            with open(filepath, 'w') as f:
                yaml.dump_all(combined_resources, f, default_flow_style=False, sort_keys=False, explicit_start=True)
            
            print(f"导出 {len(combined_resources)} 个资源到 {filepath}")
        
        # 导出没有关联Deployment的独立Service
        for svc_name, service in resources['services'].items():
            if svc_name not in exported_services:
                filename = f"{svc_name}.yaml"
                filepath = os.path.join(ns_deployments_dir, filename)
                
                with open(filepath, 'w') as f:
                    yaml.dump(service, f, default_flow_style=False, sort_keys=False)
                
                print(f"导出独立Service到 {filepath}")
    
    # 导出Ingress到ingress目录
    for namespace, resources in namespace_resources.items():
        if resources['ingresses']:
            # 在ingress目录下创建命名空间子目录
            ns_ingress_dir = os.path.join(ingress_dir, namespace)
            os.makedirs(ns_ingress_dir, exist_ok=True)
            
            # 导出所有Ingress
            for ing_name, ingress in resources['ingresses'].items():
                filename = f"{ing_name}.yaml"
                filepath = os.path.join(ns_ingress_dir, filename)
                
                with open(filepath, 'w') as f:
                    yaml.dump(ingress, f, default_flow_style=False, sort_keys=False)
                
                print(f"导出Ingress到 {filepath}")

if __name__ == "__main__":
    export_resources()